import os
import random
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

def load_model(model_path='garbage_classifier_model.keras'):
    """加载预训练模型"""
    if not os.path.exists(model_path):
        print(f"错误：模型文件 {model_path} 不存在")
        print("请确保：")
        print("1. 模型文件路径正确")
        print("2. 已经运行过训练脚本")
        return None
    
    try:
        model = tf.keras.models.load_model(model_path)
        print(f"成功加载模型: {model_path}")
        model.summary()  # 打印模型结构
        return model
    except Exception as e:
        print(f"加载模型时出错: {e}")
        return None

def load_test_images(test_dir):
    """加载测试目录中的图片"""
    if not os.path.exists(test_dir):
        print(f"错误：测试目录 {test_dir} 不存在")
        return None, None
    
    classes = ['cardboard', 'metal', 'plastic']
    images = []
    labels = []
    file_paths = []
    
    for img_file in os.listdir(test_dir):
        if img_file.lower().endswith(('.jpg', '.jpeg', '.png')):
            img_path = os.path.join(test_dir, img_file)
            try:
                # 从文件名解析标签
                label = None
                for idx, class_name in enumerate(classes):
                    if class_name.lower() in img_file.lower():
                        label = idx
                        break
                
                if label is not None:
                    img = tf.keras.preprocessing.image.load_img(
                        img_path, target_size=(128, 128))
                    img_array = tf.keras.preprocessing.image.img_to_array(img)
                    images.append(img_array)
                    labels.append(label)
                    file_paths.append(img_path)
            except Exception as e:
                print(f"加载图片 {img_file} 出错: {e}")
    
    if not images:
        print("警告：测试目录中没有有效图片")
        return None, None, None
    
    return np.array(images), np.array(labels), file_paths

def evaluate_model(model, test_dir):
    """评估模型在测试集上的表现"""
    # 加载测试数据
    X_test, y_test, file_paths = load_test_images(test_dir)
    if X_test is None:
        return
    
    # 预处理
    X_test = X_test / 255.0
    
    # 评估整体准确率
    print("\n正在评估模型...")
    test_loss, test_acc = model.evaluate(X_test, y_test, verbose=1)
    print(f"\n测试集准确率: {test_acc:.4f}, 损失: {test_loss:.4f}")
    
    # 随机选择一些样本进行可视化
    num_samples = min(5, len(X_test))
    indices = random.sample(range(len(X_test)), num_samples)
    
    plt.figure(figsize=(15, 5))
    for i, idx in enumerate(indices):
        img = X_test[idx]
        true_label = y_test[idx]
        file_path = file_paths[idx]
        
        # 预测
        pred_probs = model.predict(np.expand_dims(img, axis=0))
        pred_label = np.argmax(pred_probs)
        confidence = np.max(pred_probs)
        
        # 显示结果
        plt.subplot(1, num_samples, i+1)
        plt.imshow(img)
        plt.title(f"True: {true_label}\nPred: {pred_label}\nConf: {confidence:.2f}")
        plt.axis('off')
        
        # 打印详细信息
        print(f"\n图片: {os.path.basename(file_path)}")
        print(f"真实类别: {true_label} ({['cardboard', 'metal', 'plastic'][true_label]})")
        print(f"预测类别: {pred_label} ({['cardboard', 'metal', 'plastic'][pred_label]})")
        print(f"置信度: {confidence:.4f}")
        print("各类别概率:", {cls: f"{prob:.4f}" for cls, prob in zip(['cardboard', 'metal', 'plastic'], pred_probs[0])})
    
    plt.tight_layout()
    plt.savefig('model_predictions.png')
    plt.show()

def interactive_test(model):
    """交互式测试单个图片"""
    while True:
        print("\n交互测试模式 (输入 'q' 退出)")
        img_path = input("请输入图片路径: ").strip()
        
        if img_path.lower() == 'q':
            break
            
        if not os.path.exists(img_path):
            print("错误：文件不存在")
            continue
            
        try:
            # 预处理图片
            img = tf.keras.preprocessing.image.load_img(
                img_path, target_size=(128, 128))
            img_array = tf.keras.preprocessing.image.img_to_array(img)
            img_array = img_array / 255.0
            img_array = np.expand_dims(img_array, axis=0)
            
            # 预测
            pred_probs = model.predict(img_array)
            pred_label = np.argmax(pred_probs)
            confidence = np.max(pred_probs)
            classes = ['cardboard', 'metal', 'plastic']
            
            # 显示结果
            plt.imshow(img)
            plt.title(f"Pred: {classes[pred_label]}\nConf: {confidence:.2f}")
            plt.axis('off')
            plt.show()
            
            print("\n预测结果:")
            print(f"最可能类别: {classes[pred_label]} (置信度: {confidence:.4f})")
            print("各类别概率:")
            for cls, prob in zip(classes, pred_probs[0]):
                print(f"  {cls}: {prob:.4f}")
                
        except Exception as e:
            print(f"处理图片时出错: {e}")

if __name__ == "__main__":
    # 配置
    MODEL_PATH = 'garbage_classifier_model.keras'
    TEST_DIR = 'final_shuffled_data/test'  # 默认测试目录
    
    # 加载模型
    print("正在加载模型...")
    model = load_model(MODEL_PATH)
    if model is None:
        exit(1)
    
    # 用户选择模式
    while True:
        print("\n请选择测试模式:")
        print("1. 评估整个测试集")
        print("2. 交互式测试单张图片")
        print("3. 退出")
        choice = input("请输入选择 (1/2/3): ").strip()
        
        if choice == '1':
            # 自定义测试目录
            custom_dir = input(f"输入测试目录 (留空使用默认: {TEST_DIR}): ").strip()
            test_dir = custom_dir if custom_dir else TEST_DIR
            evaluate_model(model, test_dir)
        elif choice == '2':
            interactive_test(model)
        elif choice == '3':
            break
        else:
            print("无效输入，请重新选择")
    
    print("测试结束")